import os
import tifffile as tifi
import random
import numpy as np
from tqdm import tqdm
from Modules.ImageProcess import MyCV2 as imp
os.environ["OPENCV_IO_MAX_IMAGE_PIXELS"] = pow(2, 40).__str__()
os.environ["CV_IO_MAX_IMAGE_PIXELS"] = pow(2, 40).__str__()
import cv2
from osgeo import gdal, osr

def resize_image(img, scale_x, scale_y):
    new_width = int(img.shape[1] * scale_x)
    new_height = int(img.shape[0] * scale_y)
    resized_image = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_LINEAR)
    return resized_image


def rotate_image(img, angle):
    (height, width) = img.shape[:2]

    # 计算旋转中心（图像中心）
    center = (width // 2, height // 2)
    # 计算旋转矩阵
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, scale=1.0)

    # 计算旋转后图像的新边界尺寸，确保图像完整
    cos = np.abs(rotation_matrix[0, 0])
    sin = np.abs(rotation_matrix[0, 1])
    new_width = int((height * sin) + (width * cos))
    new_height = int((height * cos) + (width * sin))

    # 调整旋转矩阵，使旋转后的图像中心与原图中心对齐
    rotation_matrix[0, 2] += (new_width / 2) - center[0]
    rotation_matrix[1, 2] += (new_height / 2) - center[1]

    # 执行旋转并获取结果
    rotated_img = cv2.warpAffine(img, rotation_matrix, (new_width, new_height))

    return rotated_img


def shear_image(image, shear_factor_x, shear_factor_y):
    # 获取图像尺寸
    h, w = image.shape[:2]

    # 计算新的图像大小
    # 假设图像中心在变换后仍然位于中心
    # 对于x方向的shear，图像宽度会增加
    # 对于y方向的shear，图像高度会增加

    new_w = int(w + abs(shear_factor_x) * h)
    new_h = int(h + abs(shear_factor_y) * w)

    shear_matrix = np.float32([[1, shear_factor_x, 0],
                               [shear_factor_y, 1, 0],
                               [0, 0, 1]])

    # 取仿射矩阵的前两行用于warpAffine
    affine_matrix = shear_matrix[:2]
    if shear_factor_x < 0:
        affine_matrix[0, 2] = abs(shear_factor_x) * h
    if shear_factor_y < 0:
        affine_matrix[1, 2] = abs(shear_factor_y) * w

    # 应用仿射变换
    result = cv2.warpAffine(image, affine_matrix, (new_w, new_h))

    # 去除多余黑边
    gray = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)
    _, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    x, y, w, h = cv2.boundingRect(contours[0])
    result = result[y:y + h, x:x + w]
    return result


def add_gaussian_noise(image, mean=0, var=1):
    """
    向图像添加高斯噪声
    :param image: OpenCV图像对象（numpy数组）
    :param mean: 噪声的均值
    :param var: 噪声的方差
    :return: 带有噪声的OpenCV图像对象（numpy数组）
    """
    row, col, ch = image.shape
    gauss = np.random.normal(mean, var ** 0.5, (row, col, ch))
    gauss = gauss.reshape(row, col, ch)
    noisy = image + gauss

    # 确保像素值不会超出0到255的范围
    noisy = np.clip(noisy, 0, 255).astype(np.uint8)
    return noisy


def crop_img(img, crop_size):
    h, w = img.shape[:2]
    thresh = int(min(h, w) - crop_size)

    crop_h = int(crop_size)
    crop_w = int(crop_size * random.choice(scale_factor))

    coord_h, coord_w = random.randint(0, h - thresh), random.randint(0, w - thresh)
    end_h, end_w = min(coord_h + crop_h, h), min(coord_w + crop_w, w)
    croped_img = img[coord_h:end_h, coord_w:end_w, :]

    rate = (crop_h + crop_w) / 2

    return croped_img, coord_h, coord_w, end_h, end_w, rate


def get_crop_size_list(mu, sigma, num, min_size, max_size):
    random_numbers = np.random.normal(mu, sigma, num)
    random_size = list(np.clip(random_numbers, min_size, max_size))
    return random_size
def read_tif(input_file):
    # 打开原始图像
    src_ds = gdal.Open(input_file, gdal.GA_ReadOnly)
    if src_ds is None:
        print('Unable to open', input_file)
    return src_ds

def gdal_to_opencv(src_ds):
    # 打开TIFF图像
    # src_ds = gdal.Open(input_file, gdal.GA_ReadOnly)
    # if src_ds is None:
    #     print('无法打开文件:', input_file)
    #     return None

    # 获取图像的宽度和高度
    width = src_ds.RasterXSize
    height = src_ds.RasterYSize
    bands = src_ds.RasterCount

    # 根据波段数决定图像的颜色模式
    if bands == 1:  # 单波段图像（灰度）
        band = src_ds.GetRasterBand(1)
        data = band.ReadAsArray(0, 0, width, height)
        img = np.array(data, dtype=np.uint8)
    elif bands == 3:  # 三波段图像（BGR）
        b_band = src_ds.GetRasterBand(1)
        g_band = src_ds.GetRasterBand(2)
        r_band = src_ds.GetRasterBand(3)

        b = b_band.ReadAsArray(0, 0, width, height)
        g = g_band.ReadAsArray(0, 0, width, height)
        r = r_band.ReadAsArray(0, 0, width, height)

        img = np.dstack((b, g, r))  # 注意这里的顺序
        img = np.array(img, dtype=np.uint8)
        return img
    elif bands == 4:  # 四波段图像（通常为BGR+NIR或BGR+Alpha）
        b_band = src_ds.GetRasterBand(1)
        g_band = src_ds.GetRasterBand(2)
        r_band = src_ds.GetRasterBand(3)
        nir_band = src_ds.GetRasterBand(4)  # 假设第四个波段是NIR

        b = b_band.ReadAsArray(0, 0, width, height)
        g = g_band.ReadAsArray(0, 0, width, height)
        r = r_band.ReadAsArray(0, 0, width, height)
        nir = nir_band.ReadAsArray(0, 0, width, height)

        # 创建BGR图像
        img_bgr = np.dstack((b, g, r))  # 注意这里的顺序
        img_bgr = np.array(img_bgr, dtype=np.uint8)

        # 如果需要，也可以处理NIR波段
        # img_nir = np.array(nir, dtype=np.uint8)

        # 返回BGR图像
        return img_bgr
    else:
        print("不支持的波段数:", bands)
        return None

if __name__ == '__main__':
    root_path = r'G:\data\satgeoloc_finals_basemap\google_process'
    save_path = r'G:\data\satgeoloc_finals_basemap\query_image'
    os.makedirs(save_path, exist_ok=True)

    rotate_angle_list = list(np.arange(-40, -19, 10)) + list(np.arange(20, 41, 10)) + [0]
    scale_factor = list(np.round(np.arange(0.8, 1.21, 0.1), 1))

    # random.seed(0)
    image_list = os.listdir(root_path)
    for image_name in tqdm(image_list):
        # img = cv2.imread(f'{root_path}/{image_name}',cv2.IMREAD_UNCHANGED)
        # desired_order = [2, 1, 0]
        # images = [tifi.imread(f'{root_path}/{image_name}', key=i) for i in desired_order]
        # stacked_images = np.stack(images, axis=0)
        img = tifi.imread(f'{root_path}/{image_name}')[:, :, [2, 1, 0]]
        # img = imp.cv_imread(f'{root_path}/{image_name}')
        # img = read_tif(f'{root_path}/{image_name}')
        # img = gdal_to_opencv(img)
        if (img.shape[0] > 24000) & (img.shape[1] > 24000):
            num = 900
        else:
            num = 700
        random_size = get_crop_size_list(5000, 2000, num, 300, 10000)

        for crop_size in random_size:
            croped_img, coord_h, coord_w, end_h, end_w, rate = crop_img(img, crop_size)
            result = cv2.resize(croped_img, (448, 448))

            rotate_angle = random.choice(rotate_angle_list)
            scale_x, scale_y = random.choice(scale_factor), random.choice(scale_factor)

            result = resize_image(result, scale_x, scale_y)
            result = rotate_image(result, rotate_angle)  # 旋转
            result = add_gaussian_noise(result)

            cv2.imwrite(f'{save_path}/{image_name[:-4]}_{round(rate/448, 2)}_{coord_w}_{coord_h}_{end_w}_{end_h}.jpg', result)
    print

